/**
 * Project Name: socket-test
 * Package Name: com.meatball
 * File Name: RoomHandler
 * Date: 2021/3/8 17:19
 * User: 張翔宇
 * Copyright (c) 2021, zhang.xiangyu@foxmail.com All Rights Reserved.
 */
package com.meatball;

import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.HandshakeInfo;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 聊天室
 * @title: RoomHandler.java
 * @author: 張翔宇
 * @date: 2021/3/8 17:19
 * @version: V3.0
 */
@Slf4j
@Component
@WebSocketMapping("/room")
public class RoomHandler implements WebSocketHandler {
    public static ConcurrentHashMap<String, Map<String, WebSocketSender>> roomCacheMap = new ConcurrentHashMap<>();

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        HandshakeInfo handshakeInfo = session.getHandshakeInfo();
        InetSocketAddress address = handshakeInfo.getRemoteAddress();
        String params = handshakeInfo.getUri().getQuery();

        Map<String, String> paramMap = HttpUtil.decodeParamMap(params, StandardCharsets.UTF_8);

        String roomName = paramMap.get("roomName");
        String userId = paramMap.get("userId");

        // 出站
        Mono<Void> output = session.send(Flux.create(sink -> handleClient(
                roomName,
                userId,
                new WebSocketSender(session, sink)
        )));

        // 入站
        Mono<Void> input = session.receive()
                .doOnSubscribe(conn -> log.info("建立链接：{}，用户IP：{}，房间号：{}，用户：{}"
                        , session.getId(), address.getHostName(), roomName, userId))
                .doOnNext(msg -> broadcast(roomName, userId, msg.getPayloadAsText()))
                .doOnComplete(() -> {
                    log.info("关闭连接：{}", session.getId());
                    exitRoom(session, roomName, userId);
                })
                .doOnCancel(() -> {
                    log.info("关闭连接：{}", session.getId());
                    exitRoom(session, roomName, userId);
                })
                .then();

        return Mono.zip(input, output).then();
    }

    private void exitRoom(WebSocketSession session, String roomName, String userId) {
        session.close().share().then();
        broadcast(roomName, userId, "退出房间！");
        removeUser(roomName, userId);
    }

    private void removeUser(String roomName, String userId) {
        log.info("用户：{}，退出房间：{}", userId, roomName);
        Map<String, WebSocketSender> socketClientCacheMap = roomCacheMap.get(roomName);
        socketClientCacheMap.remove(userId);
        if (socketClientCacheMap.isEmpty()) {
            log.info("房间：{}没人了，关闭房间！", roomName);
            roomCacheMap.remove(roomName);
        }
    }

    private void handleClient(String roomName, String userId, WebSocketSender sender) {
        if (!roomCacheMap.containsKey(roomName)) {
            log.info("用户：{}，创建房间：{}", userId, roomName);
            Map<String, WebSocketSender> socketClientCacheMap = new HashMap<>();
            socketClientCacheMap.put(userId, sender);
            roomCacheMap.put(roomName, socketClientCacheMap);
        } else {
            Map<String, WebSocketSender> socketClientCacheMap = roomCacheMap.get(roomName);
            if (!socketClientCacheMap.containsKey(userId)) {
                log.info("用户：{}，进入房间：{}", userId, roomName);
                socketClientCacheMap.put(userId, sender);
            }
        }
    }

    /**
     * 发布广播
     * @title: broadcast
     * @param roomName
     * @param userId
     * @param message
     * @return: void 返回类型
     * @author: 張翔宇
     * @date: 2021/3/8 17:48
     */
    public void broadcast(String roomName, String userId, String message) {
        JSONObject msgObj = new JSONObject();
        msgObj.put("roomName", roomName);
        msgObj.put("userId", userId);
        msgObj.put("message", message);

        sendToAll(roomName, userId, message);
    }

    /**
     * 发布消息给除了自己的所有用户
     * @title: sendToAll
     * @param roomName
     * @param userId
     * @param message
     * @return: void 返回类型
     * @author: 張翔宇
     * @date: 2021/3/8 17:49
     */
    public static void sendToAll(String roomName, String userId, String message) {
        Map<String, WebSocketSender> senders = roomCacheMap.get(roomName);
        senders.forEach((user, sender) -> {
            if (!userId.equals(user)) {
                log.info("用户：{}发送消息：{}", userId, message);
                sender.sendData(userId + "：" + message);
            }
        });
    }
}
